package io.kwdb.sdk.internal;

import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.*;

/**
 * Simple metrics collector for monitoring SDK performance. Provides counters, gauges, and latency
 * histograms.
 */
public class Metrics {

  // Counters
  private final AtomicLong writesOk = new AtomicLong(0);
  private final Map<Integer, AtomicLong> writesFailed = new ConcurrentHashMap<>();
  private final AtomicLong droppedItems = new AtomicLong(0);
  private final AtomicLong retries = new AtomicLong(0);

  // Gauges
  private final AtomicInteger currentBatchSize = new AtomicInteger(0);
  private final AtomicInteger currentQueueSize = new AtomicInteger(0);

  // Latency tracking (simplified percentile calculation)
  private final ConcurrentSkipListMap<Long, AtomicLong> batchLatencies =
      new ConcurrentSkipListMap<>();
  private final ConcurrentSkipListMap<Long, AtomicLong> writeLatencies =
      new ConcurrentSkipListMap<>();
  private final Object latencyLock = new Object();

  /** Increment successful writes counter. */
  public void incrementWritesOk(long count) {
    writesOk.addAndGet(count);
  }

  /** Increment failed writes counter by error code. */
  public void incrementWritesFailed(long count, int errorCode) {
    writesFailed.computeIfAbsent(errorCode, k -> new AtomicLong(0)).addAndGet(count);
  }

  /** Record batch size. */
  public void observeBatchSize(int size) {
    currentBatchSize.set(size);
  }

  /** Record batch processing latency. */
  public void observeBatchLatency(long latencyMs) {
    recordLatency(batchLatencies, latencyMs);
  }

  /** Record individual write latency. */
  public void observeWriteLatency(long latencyMs) {
    recordLatency(writeLatencies, latencyMs);
  }

  /** Increment dropped items counter. */
  public void incrementDropped(long count) {
    droppedItems.addAndGet(count);
  }

  /** Increment retry counter. */
  public void incrementRetries() {
    retries.incrementAndGet();
  }

  /** Update current queue size. */
  public void updateQueueSize(int size) {
    currentQueueSize.set(size);
  }

  /** Get a snapshot of current metrics. */
  public MetricsSnapshot getSnapshot() {
    Map<Integer, Long> failedCopy = new HashMap<>();
    writesFailed.forEach((k, v) -> failedCopy.put(k, v.get()));

    return new MetricsSnapshot(
        writesOk.get(),
        failedCopy,
        droppedItems.get(),
        retries.get(),
        currentBatchSize.get(),
        currentQueueSize.get(),
        calculatePercentiles(batchLatencies),
        calculatePercentiles(writeLatencies));
  }

  /** Reset all metrics. */
  public void reset() {
    writesOk.set(0);
    writesFailed.clear();
    droppedItems.set(0);
    retries.set(0);
    currentBatchSize.set(0);
    currentQueueSize.set(0);
    batchLatencies.clear();
    writeLatencies.clear();
  }

  /** Record a latency sample. */
  private void recordLatency(ConcurrentSkipListMap<Long, AtomicLong> map, long latencyMs) {
    synchronized (latencyLock) {
      map.computeIfAbsent(latencyMs, k -> new AtomicLong(0)).incrementAndGet();

      // Keep only recent 1000 samples to avoid memory growth
      while (map.size() > 1000) {
        map.pollFirstEntry();
      }
    }
  }

  /** Calculate percentiles from latency samples. */
  private LatencyPercentiles calculatePercentiles(ConcurrentSkipListMap<Long, AtomicLong> samples) {
    if (samples.isEmpty()) {
      return new LatencyPercentiles(0, 0, 0, 0, 0);
    }

    List<Long> values = new ArrayList<>();
    samples.forEach(
        (latency, count) -> {
          for (int i = 0; i < count.get(); i++) {
            values.add(latency);
          }
        });

    if (values.isEmpty()) {
      return new LatencyPercentiles(0, 0, 0, 0, 0);
    }

    Collections.sort(values);
    int size = values.size();

    return new LatencyPercentiles(
        values.get(0), // min
        values.get((int) (size * 0.5)), // p50
        values.get((int) (size * 0.95)), // p95
        values.get((int) (size * 0.99)), // p99
        values.get(size - 1) // max
        );
  }

  /** Snapshot of metrics at a point in time. */
  public static class MetricsSnapshot {
    public final long writesOk;
    public final Map<Integer, Long> writesFailed;
    public final long droppedItems;
    public final long retries;
    public final int currentBatchSize;
    public final int currentQueueSize;
    public final LatencyPercentiles batchLatency;
    public final LatencyPercentiles writeLatency;

    public MetricsSnapshot(
        long writesOk,
        Map<Integer, Long> writesFailed,
        long droppedItems,
        long retries,
        int currentBatchSize,
        int currentQueueSize,
        LatencyPercentiles batchLatency,
        LatencyPercentiles writeLatency) {
      this.writesOk = writesOk;
      this.writesFailed = Collections.unmodifiableMap(writesFailed);
      this.droppedItems = droppedItems;
      this.retries = retries;
      this.currentBatchSize = currentBatchSize;
      this.currentQueueSize = currentQueueSize;
      this.batchLatency = batchLatency;
      this.writeLatency = writeLatency;
    }

    @Override
    public String toString() {
      StringBuilder sb = new StringBuilder();
      sb.append("Metrics{\n");
      sb.append("  writes: ok=").append(writesOk);
      if (!writesFailed.isEmpty()) {
        sb.append(", failed=").append(writesFailed);
      }
      sb.append("\n");
      sb.append("  dropped: ").append(droppedItems).append("\n");
      sb.append("  retries: ").append(retries).append("\n");
      sb.append("  batch: size=")
          .append(currentBatchSize)
          .append(", queueSize=")
          .append(currentQueueSize)
          .append("\n");
      sb.append("  batchLatency: ").append(batchLatency).append("\n");
      sb.append("  writeLatency: ").append(writeLatency).append("\n");
      sb.append("}");
      return sb.toString();
    }
  }

  /** Latency percentile statistics. */
  public static class LatencyPercentiles {
    public final long min;
    public final long p50;
    public final long p95;
    public final long p99;
    public final long max;

    public LatencyPercentiles(long min, long p50, long p95, long p99, long max) {
      this.min = min;
      this.p50 = p50;
      this.p95 = p95;
      this.p99 = p99;
      this.max = max;
    }

    @Override
    public String toString() {
      return String.format(
          "min=%dms, p50=%dms, p95=%dms, p99=%dms, max=%dms", min, p50, p95, p99, max);
    }
  }

  /** Global metrics instance. */
  private static final Metrics INSTANCE = new Metrics();

  /** Get the global metrics instance. */
  public static Metrics getInstance() {
    return INSTANCE;
  }
}
